#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
视频处理器模块
提供视频信息提取、帧采样和视频分段等功能
"""

import os
import cv2
import numpy as np
from typing import Dict, Any, List, Optional, Tuple
import hashlib
from datetime import datetime


class VideoProcessor:
    """
    视频处理器
    提供视频基本信息提取、关键帧采样、视频分段等功能
    """
    
    def __init__(self):
        """
        初始化视频处理器
        """
        # 视频处理配置
        self.default_max_frames = 5  # 默认最大采样帧数
        self.default_frame_interval = 3  # 默认帧间隔（秒）
        self.min_frame_quality = 0.3  # 最低帧质量阈值
        self.temp_frame_dir = None  # 临时帧保存目录
    
    def set_temp_directory(self, temp_dir: str):
        """
        设置临时帧保存目录
        
        Args:
            temp_dir: 临时目录路径
        """
        self.temp_frame_dir = temp_dir
        # 确保目录存在
        if temp_dir and not os.path.exists(temp_dir):
            os.makedirs(temp_dir)
    
    def extract_video_info(self, video_path: str) -> Dict[str, Any]:
        """
        提取视频基本信息
        
        Args:
            video_path: 视频文件路径
            
        Returns:
            Dict[str, Any]: 视频信息字典
        """
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"视频文件不存在: {video_path}")
        
        # 打开视频文件
        cap = cv2.VideoCapture(video_path)
        
        if not cap.isOpened():
            raise ValueError(f"无法打开视频文件: {video_path}")
        
        try:
            # 获取视频信息
            fps = cap.get(cv2.CAP_PROP_FPS)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            duration = total_frames / fps if fps > 0 else 0
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            codec = int(cap.get(cv2.CAP_PROP_FOURCC))
            
            # 计算视频大小
            file_size = os.path.getsize(video_path) / (1024 * 1024)  # MB
            
            # 获取文件名和扩展名
            file_name = os.path.basename(video_path)
            file_ext = os.path.splitext(file_name)[1].lower()
            
            # 生成视频唯一标识符
            video_id = self._generate_video_id(video_path)
            
            # 视频信息字典
            video_info = {
                "file_path": video_path,
                "file_name": file_name,
                "file_extension": file_ext,
                "file_size_mb": round(file_size, 2),
                "video_id": video_id,
                "duration_seconds": round(duration, 2),
                "fps": round(fps, 2),
                "total_frames": total_frames,
                "width": width,
                "height": height,
                "resolution": f"{width}x{height}",
                "codec": codec,
                "codec_name": self._fourcc_to_string(codec),
                "aspect_ratio": round(width / height, 2) if height > 0 else 0,
                "creation_time": self._get_file_creation_time(video_path),
                "last_modified_time": self._get_file_modified_time(video_path)
            }
            
            # 计算视频特征
            video_info["video_category"] = self._classify_video_category(video_info)
            
            return video_info
            
        finally:
            # 释放资源
            cap.release()
    
    def extract_key_frames(self, video_path: str, 
                          max_frames: int = None,
                          frame_interval: int = None,
                          strategy: str = "uniform",
                          save_frames: bool = False) -> List[Dict[str, Any]]:
        """
        提取视频关键帧
        
        Args:
            video_path: 视频文件路径
            max_frames: 最大帧数，None表示使用默认值
            frame_interval: 帧间隔（秒），None表示使用默认值
            strategy: 采样策略 - "uniform"(均匀)或"content_based"(基于内容)
            save_frames: 是否保存帧到临时目录
            
        Returns:
            List[Dict[str, Any]]: 关键帧信息列表，包含帧数据和元数据
        """
        if max_frames is None:
            max_frames = self.default_max_frames
        
        if frame_interval is None:
            frame_interval = self.default_frame_interval
        
        # 获取视频信息
        video_info = self.extract_video_info(video_path)
        fps = video_info["fps"]
        total_frames = video_info["total_frames"]
        duration = video_info["duration_seconds"]
        
        # 打开视频文件
        cap = cv2.VideoCapture(video_path)
        
        if not cap.isOpened():
            raise ValueError(f"无法打开视频文件: {video_path}")
        
        try:
            # 计算采样帧位置
            sample_positions = self._calculate_sample_positions(
                total_frames, fps, duration, max_frames, frame_interval, strategy
            )
            
            # 提取帧
            key_frames = []
            prev_frame = None
            
            for frame_idx in sample_positions:
                # 设置帧位置
                cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
                
                # 读取帧
                ret, frame = cap.read()
                if ret:
                    # 计算帧质量指标
                    frame_time = frame_idx / fps if fps > 0 else 0
                    frame_info = {
                        "frame_index": frame_idx,
                        "frame_time": round(frame_time, 2),
                        "quality_score": self._calculate_frame_quality(frame, prev_frame),
                        "width": frame.shape[1],
                        "height": frame.shape[0]
                    }
                    
                    # 如果需要保存帧
                    if save_frames and self.temp_frame_dir:
                        frame_path = self._save_frame(frame, video_info["video_id"], frame_idx)
                        frame_info["frame_path"] = frame_path
                    
                    # 保存帧数据
                    _, buffer = cv2.imencode('.jpg', frame)
                    frame_info["frame_bytes"] = buffer.tobytes()
                    
                    key_frames.append(frame_info)
                    prev_frame = frame
            
            # 对帧进行质量排序（如果是基于内容的采样）
            if strategy == "content_based":
                key_frames.sort(key=lambda x: x["quality_score"], reverse=True)
                key_frames = key_frames[:max_frames]  # 只保留最高分的帧
                
                # 按时间顺序排序
                key_frames.sort(key=lambda x: x["frame_time"])
            
            print(f"成功提取 {len(key_frames)} 个关键帧")
            return key_frames
            
        finally:
            cap.release()
    
    def segment_video(self, video_path: str, segment_length: int = 30) -> List[Dict[str, Any]]:
        """
        将视频分段
        
        Args:
            video_path: 视频文件路径
            segment_length: 分段长度（秒）
            
        Returns:
            List[Dict[str, Any]]: 视频分段信息列表
        """
        # 获取视频信息
        video_info = self.extract_video_info(video_path)
        duration = video_info["duration_seconds"]
        fps = video_info["fps"]
        
        # 计算分段数量
        segment_count = int(np.ceil(duration / segment_length))
        
        # 生成分段信息
        segments = []
        for i in range(segment_count):
            start_time = i * segment_length
            end_time = min((i + 1) * segment_length, duration)
            
            # 计算起始和结束帧数
            start_frame = int(start_time * fps)
            end_frame = int(end_time * fps)
            
            segment = {
                "segment_id": f"{video_info['video_id']}_seg{i+1:03d}",
                "start_time": round(start_time, 2),
                "end_time": round(end_time, 2),
                "duration": round(end_time - start_time, 2),
                "start_frame": start_frame,
                "end_frame": end_frame,
                "frame_count": end_frame - start_frame + 1,
                "segment_index": i + 1,
                "total_segments": segment_count,
                "video_id": video_info["video_id"]
            }
            
            segments.append(segment)
        
        return segments
    
    def analyze_video_composition(self, video_path: str) -> Dict[str, Any]:
        """
        分析视频构图特征
        
        Args:
            video_path: 视频文件路径
            
        Returns:
            Dict[str, Any]: 视频构图分析结果
        """
        # 提取少量关键帧进行分析
        key_frames = self.extract_key_frames(video_path, max_frames=10, strategy="content_based")
        
        if not key_frames:
            return {"error": "无法提取关键帧进行分析"}
        
        # 分析各帧的特征
        brightness_values = []
        color_distribution = []
        edge_density = []
        motion_intensity = []
        
        prev_frame = None
        
        for i, frame_info in enumerate(key_frames):
            # 读取帧数据
            frame_bytes = frame_info["frame_bytes"]
            frame = cv2.imdecode(np.frombuffer(frame_bytes, np.uint8), cv2.IMREAD_COLOR)
            
            # 转换为灰度图
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 计算亮度
            brightness = np.mean(gray) / 255.0
            brightness_values.append(brightness)
            
            # 计算色彩分布
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            hist_h = cv2.calcHist([hsv], [0], None, [8], [0, 180])
            color_distribution.append(np.mean(hist_h).tolist())
            
            # 计算边缘密度
            edges = cv2.Canny(gray, 100, 200)
            density = np.sum(edges > 0) / (frame.shape[0] * frame.shape[1])
            edge_density.append(density)
            
            # 计算运动强度
            if prev_frame is not None:
                prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
                flow = cv2.calcOpticalFlowFarneback(
                    prev_gray, gray, None, 0.5, 3, 15, 3, 5, 1.2, 0
                )
                magnitude = np.sqrt(flow[..., 0]**2 + flow[..., 1]**2)
                motion = np.mean(magnitude)
                motion_intensity.append(motion)
            
            prev_frame = frame.copy()
        
        # 计算平均特征
        composition_analysis = {
            "brightness": {
                "average": round(np.mean(brightness_values), 3),
                "std": round(np.std(brightness_values), 3),
                "max": round(np.max(brightness_values), 3),
                "min": round(np.min(brightness_values), 3)
            },
            "color_diversity": {
                "average": round(np.mean(color_distribution), 3),
                "std": round(np.std(color_distribution), 3)
            },
            "edge_density": {
                "average": round(np.mean(edge_density), 3),
                "std": round(np.std(edge_density), 3)
            },
            "frame_count_analyzed": len(key_frames),
            "video_style": self._infer_video_style(brightness_values, edge_density)
        }
        
        # 添加运动强度分析（如果有）
        if motion_intensity:
            composition_analysis["motion_intensity"] = {
                "average": round(np.mean(motion_intensity), 3),
                "std": round(np.std(motion_intensity), 3)
            }
        
        return composition_analysis
    
    def _generate_video_id(self, video_path: str) -> str:
        """
        生成视频唯一标识符
        
        Args:
            video_path: 视频文件路径
            
        Returns:
            str: 视频ID
        """
        # 使用文件路径和修改时间生成唯一ID
        file_stat = os.stat(video_path)
        unique_str = f"{video_path}_{file_stat.st_mtime}_{file_stat.st_size}"
        
        # 计算MD5
        md5_hash = hashlib.md5(unique_str.encode()).hexdigest()
        
        # 返回前16位作为视频ID
        return f"video_{md5_hash[:16]}"
    
    def _fourcc_to_string(self, fourcc: int) -> str:
        """
        将FOURCC代码转换为字符串
        
        Args:
            fourcc: FOURCC代码
            
        Returns:
            str: 编解码器名称
        """
        try:
            return "".join([chr((fourcc >> 8 * i) & 0xFF) for i in range(4)])
        except:
            return "unknown"
    
    def _get_file_creation_time(self, file_path: str) -> str:
        """
        获取文件创建时间
        
        Args:
            file_path: 文件路径
            
        Returns:
            str: 格式化的时间字符串
        """
        file_stat = os.stat(file_path)
        try:
            # 尝试获取创建时间
            if hasattr(file_stat, 'st_birthtime'):
                # macOS
                timestamp = file_stat.st_birthtime
            else:
                # Windows或Linux
                timestamp = file_stat.st_ctime
            return datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:%M:%S")
        except:
            return "unknown"
    
    def _get_file_modified_time(self, file_path: str) -> str:
        """
        获取文件修改时间
        
        Args:
            file_path: 文件路径
            
        Returns:
            str: 格式化的时间字符串
        """
        file_stat = os.stat(file_path)
        return datetime.fromtimestamp(file_stat.st_mtime).strftime("%Y-%m-%d %H:%M:%S")
    
    def _classify_video_category(self, video_info: Dict[str, Any]) -> str:
        """
        根据视频信息分类视频类型
        
        Args:
            video_info: 视频信息
            
        Returns:
            str: 视频类别
        """
        duration = video_info["duration_seconds"]
        resolution = video_info["resolution"]
        aspect_ratio = video_info["aspect_ratio"]
        
        # 根据时长分类
        if duration < 60:  # 小于1分钟
            if aspect_ratio > 1.7:
                return "短视频(竖屏)"
            else:
                return "短视频(横屏)"
        elif duration < 300:  # 1-5分钟
            return "中短视频"
        elif duration < 1800:  # 5-30分钟
            return "中等长度视频"
        else:
            return "长视频"
    
    def _calculate_sample_positions(self, total_frames: int, fps: float, 
                                  duration: float, max_frames: int,
                                  frame_interval: int, strategy: str) -> List[int]:
        """
        计算采样帧位置
        
        Args:
            total_frames: 总帧数
            fps: 帧率
            duration: 时长
            max_frames: 最大帧数
            frame_interval: 帧间隔（秒）
            strategy: 采样策略
            
        Returns:
            List[int]: 采样帧索引列表
        """
        if strategy == "content_based":
            # 基于内容的采样，先采集更多的帧，然后根据质量筛选
            sample_count = min(total_frames, max_frames * 3)
            step = max(1, total_frames // sample_count)
            return [i * step for i in range(sample_count)]
        else:
            # 均匀采样
            if duration <= frame_interval * (max_frames - 1):
                # 视频较短，采样所有可能的帧
                frame_interval_frames = int(fps * frame_interval)
                if frame_interval_frames < 1:
                    frame_interval_frames = 1
                
                sample_positions = []
                for i in range(0, total_frames, frame_interval_frames):
                    sample_positions.append(i)
            else:
                # 视频较长，均匀采样
                step = total_frames / (max_frames + 1)
                sample_positions = [int(step * i) for i in range(1, max_frames + 1)]
            
            # 确保采样帧数不超过最大限制
            return sample_positions[:max_frames]
    
    def _calculate_frame_quality(self, frame: np.ndarray, 
                               prev_frame: Optional[np.ndarray] = None) -> float:
        """
        计算帧质量分数
        
        Args:
            frame: 当前帧
            prev_frame: 前一帧
            
        Returns:
            float: 质量分数（0-1）
        """
        # 计算亮度
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        brightness = np.mean(gray) / 255.0
        
        # 计算对比度
        contrast = np.std(gray) / 255.0
        
        # 计算清晰度（使用拉普拉斯算子）
        laplacian = cv2.Laplacian(gray, cv2.CV_64F)
        sharpness = np.var(laplacian) / 10000  # 归一化
        
        # 计算色彩丰富度
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        hue_range = np.max(hsv[:, :, 0]) - np.min(hsv[:, :, 0])
        color_diversity = min(hue_range / 180.0, 1.0)
        
        # 计算与前一帧的差异（如果有）
        motion_score = 0.3  # 默认值
        if prev_frame is not None:
            prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
            diff = cv2.absdiff(gray, prev_gray)
            motion_score = min(np.mean(diff) / 50.0, 1.0)  # 归一化
        
        # 计算最终质量分数（加权平均）
        # 亮度适中（0.3-0.7）给高分
        brightness_score = 1.0 - abs(brightness - 0.5) * 2
        
        # 综合评分
        quality_score = (
            brightness_score * 0.2 +
            contrast * 0.2 +
            sharpness * 0.2 +
            color_diversity * 0.2 +
            motion_score * 0.2
        )
        
        return min(quality_score, 1.0)
    
    def _save_frame(self, frame: np.ndarray, video_id: str, 
                   frame_idx: int) -> str:
        """
        保存帧到临时目录
        
        Args:
            frame: 帧数据
            video_id: 视频ID
            frame_idx: 帧索引
            
        Returns:
            str: 保存的文件路径
        """
        if not self.temp_frame_dir:
            raise ValueError("临时目录未设置")
        
        # 创建视频专用子目录
        video_frame_dir = os.path.join(self.temp_frame_dir, video_id)
        if not os.path.exists(video_frame_dir):
            os.makedirs(video_frame_dir)
        
        # 生成文件名
        frame_filename = f"frame_{frame_idx:06d}.jpg"
        frame_path = os.path.join(video_frame_dir, frame_filename)
        
        # 保存帧
        cv2.imwrite(frame_path, frame)
        
        return frame_path
    
    def _infer_video_style(self, brightness_values: List[float], 
                         edge_density: List[float]) -> str:
        """
        根据视频特征推断视频风格
        
        Args:
            brightness_values: 亮度值列表
            edge_density: 边缘密度列表
            
        Returns:
            str: 视频风格描述
        """
        avg_brightness = np.mean(brightness_values)
        avg_edge_density = np.mean(edge_density)
        
        # 判断风格
        if avg_brightness > 0.6 and avg_edge_density < 0.1:
            return "明亮简约风格"
        elif avg_brightness > 0.6 and avg_edge_density > 0.15:
            return "明亮细节丰富风格"
        elif avg_brightness < 0.4 and avg_edge_density < 0.1:
            return "暗色调简约风格"
        elif avg_brightness < 0.4 and avg_edge_density > 0.15:
            return "暗色调细节丰富风格"
        elif avg_edge_density > 0.2:
            return "高对比度风格"
        else:
            return "均衡风格"


# 使用示例
if __name__ == '__main__':
    try:
        # 创建视频处理器实例
        processor = VideoProcessor()
        
        # 设置临时目录
        processor.set_temp_directory("../temp/frames")
        
        # 示例视频路径（需要替换为实际路径）
        # video_path = "../data/MP4/sample_video.mp4"
        
        # 如果有示例视频，可以取消注释以下代码进行测试
        # print("提取视频信息...")
        # video_info = processor.extract_video_info(video_path)
        # print("视频信息:")
        # for key, value in video_info.items():
        #     print(f"  {key}: {value}")
        # 
        # print("\n提取关键帧...")
        # key_frames = processor.extract_key_frames(video_path, strategy="content_based")
        # print(f"成功提取 {len(key_frames)} 个关键帧")
        # 
        # print("\n视频分段...")
        # segments = processor.segment_video(video_path, segment_length=15)
        # print(f"视频分为 {len(segments)} 段")
        # 
        # print("\n分析视频构图...")
        # composition = processor.analyze_video_composition(video_path)
        # print("构图分析结果:")
        # print(f"  视频风格: {composition.get('video_style')}")
        # print(f"  平均亮度: {composition.get('brightness', {}).get('average')}")
        # print(f"  边缘密度: {composition.get('edge_density', {}).get('average')}")
        
        print("视频处理器初始化成功!")
        print("请在实际使用时，提供有效的视频路径进行处理。")
        
    except Exception as e:
        print(f"初始化失败: {str(e)}")